# -*- coding: utf-8 -*-
import re
import os
import requests
import json
import zipfile
import shutil

class DriverSpoderCLass:

    def __init__(self, save_dir="", proxy=None) -> None:
        """
        初始化driver爬虫
        """
        self.save_dir = save_dir
        self.proxy = proxy
        self.ua = 'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/94.0.4606.71 Safari/537.36 Core/1.94.249.400 QQBrowser/12.5.5658.400'
        self.headers = {
            'User-Agent': self.ua
        }
        self.googlechromelabs_url = "http://googlechromelabs.github.io/chrome-for-testing/known-good-versions-with-downloads.json"

    def get_googlechromelabs_data(self):
        """
        获取driver信息
        """
        try:
            req = requests.get(self.googlechromelabs_url, timeout=60, verify=False, headers=self.headers, proxies=self.proxy)
        except Exception as e:
            print(str(e))
            return None
        try:
            data = json.loads(req.text)
        except Exception as e:
            print("页面数据解析失败")
            return None
        return data

    def fetch_driver_url_map(self, basic_data_map: dict, platform: str="win32", version: int=0):
        """
        获取driver下载路径
        """
        data_map = {}
        basic_data_map = basic_data_map["versions"]
        for _basic_data in basic_data_map[::-1]:
            _version = re.findall("^(\d+)\.", _basic_data["version"])[0]
            if version and str(version) != _version:
                continue
            if _version in data_map:
                continue
            try:
                url_map = _basic_data["downloads"]["chromedriver"]
            except:
                continue
            for url_data in url_map:
                if url_data["platform"] == platform:
                    data_map[_version] = url_data["url"]
                    break
        return data_map
    
    def dumps_driver_files(self, data_map: dict):
        """
        缓存driver文件
        """
        driverfile_in_zip = r"chromedriver-win32/chromedriver.exe"
        basic_dir = os.path.join(os.getcwd(), self.save_dir, 'drivers')
        if not os.path.isdir(basic_dir):
            os.mkdir(basic_dir)
        for version, url in data_map.items():
            dir_name = os.path.join(basic_dir, version)
            zip_name = os.path.split(url)[-1]
            driver_path = os.path.join(dir_name, 'chromedriver.exe')
            zip_filename = os.path.join(dir_name, zip_name)
            zip_filepath = os.path.join(dir_name, zip_filename)
            if not os.path.isdir(dir_name):
                os.mkdir(dir_name)
            if os.path.isfile(driver_path):
                continue
            try:
                res = requests.get(url, timeout=20, proxies=self.proxy).content
            except Exception as e:
                print("无法下载对应driver文件，请检测网络")
                return None
            with open(zip_filepath, "wb") as f:
                f.write(res)
                f.close()
            # 解压zip
            z = zipfile.ZipFile(zip_filepath, "r")
            try:
                z.extract(driverfile_in_zip, dir_name)
            except Exception as e:
                raise e
            z.close()
            # 迁移
            oldfile = os.path.abspath(os.path.join(dir_name, driverfile_in_zip))
            shutil.move(oldfile, driver_path)
            os.remove(zip_filepath)
            os.removedirs(os.path.split(oldfile)[0])
            print(f'{version} download driver OK')
        return True

        


def demo():
    proxy = {
        'http': 'http://127.0.0.1:10809',
        'https': 'http://127.0.0.1:10809',
    }
    proxy = None
    D = DriverSpoderCLass(proxy=proxy)
    data = D.get_googlechromelabs_data()
    url_map = D.fetch_driver_url_map(data)
    print(url_map)
    D.dumps_driver_files(url_map)


if __name__ == "__main__":
    demo()
